import csv
import sqlite3
import pandas as pd
import utils.configFinRAG as configFinRAG

def query_db(sql, cursor):
    """执行 SQL 查询并返回执行结果和成功标志"""
    try:
        cursor.execute(sql)
        result_cols = cursor.fetchall()
        exc_result = str(result_cols)
        success_flag = 1
    except sqlite3.Error as e:
        exc_result = f"error: {e}"
        success_flag = 0
        print(f"SQL execution error: {e}")
    return success_flag, exc_result

def process_sql_file(sql_file_path, output_file_path, db_path):
    """处理 SQL 文件并写入执行结果"""
    # 读取 SQL 文件
    question_sql_file = pd.read_csv(sql_file_path, delimiter=",", header=0)

    # 打开 SQL 执行结果文件
    with open(output_file_path, 'w', newline='', encoding='utf-8-sig') as file:
        csvwriter = csv.writer(file)
        csvwriter.writerow(['问题id', '问题', 'SQL', 'flag', '执行结果'])

        # 数据库连接
        with sqlite3.connect(db_path) as conn:
            cursor = conn.cursor()

            # 执行 SQL 并返回结果写入 SQL 执行结果文件
            for index, row in question_sql_file.iterrows():
                if row['SQL'] != 'error':
                    success_flag, exc_result = query_db(row['SQL'], cursor)
                    csvwriter.writerow([
                        str(row['问题id']),
                        str(row['问题']),
                        row['SQL'],
                        success_flag,
                        exc_result
                    ])

if __name__ == '__main__':
    # 配置文件路径
    db_path = r'C:\Users\langzi\Desktop\local_human_number\数字人\Fay-fay-agent-edition0830\agent\tools\data\博金杯比赛数据.db'

    # 处理 SQL 文件
    process_sql_file(configFinRAG.question_sql_path, configFinRAG.question_sql_check_path, db_path)